import os
import strutils
import tables

var map = initTable[string, bool]()

proc badDir(dirName: string): bool = 
  if map.hasKey(dirName):
    return map[dirName]
  const skipDirectories = ["target", "build", "node_modules", "dist", ".git"]
  for skipDir in skipDirectories:
    if dirName == skipDir or dirName == "/" & skipDir or dirName == "\\" & skipDir:
      map[dirName] = true 
      return true
  map[dirName] = false 
  return false 

proc shouldSkipDirectory(directoryName: string): bool =
  var tmpPath = directoryName 
  if badDir(tmpPath.splitPath.tail):
    return true 
  while tmpPath.splitPath.head != "" and tmpPath.splitPath.head != "/" and tmpPath.splitPath.head != "\\": 
    if badDir(tmpPath.splitPath.tail):
      return true 
    tmpPath = tmpPath.splitPath.head
    if badDir(tmpPath.splitPath.tail):
      return true  
    
  return false

proc countLinesInDirectory(directory: string, fileExtension: string): int =
  var totalLines = 0
  var paths: seq[string] = @[]
  for entry in walkDirRec(directory):
    let filePath = entry
    let head = filePath.splitPath.head
    let dirName: string = head[len(getCurrentDir())..<len(head)]
    if shouldSkipDirectory(dirName) == false and filePath.endsWith(fileExtension):
      paths.add(filePath)
  for filePath in paths:
    var file = open(filePath)
    var buf = ""
    while file.readLine(buf):
      totalLines += 1
    file.close()
  return totalLines

when isMainModule:
  let currentDirectory = getCurrentDir()
  let fileExtension = "." & paramStr(1)
  let totalLines = countLinesInDirectory(currentDirectory, fileExtension)
  echo "Total lines of code in files with extension ", fileExtension, ": ", totalLines
